#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue

class KernelFastGeluGrad {
public:
    __aicore__ inline KernelFastGeluGrad() {}
    __aicore__ inline void Init(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices,
                                uint32_t CoreDataNum, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t TailDataNum,
                                uint32_t xyz2DataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->coreDataNum = CoreDataNum;
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = TailDataNum;
        this->xyz2DataNum =  xyz2DataNum;

        xyz1Gm.SetGlobalBuffer((__gm__ DTYPE_XYZ1*)xyz1, this->coreDataNum);
        xyz2Gm.SetGlobalBuffer((__gm__ DTYPE_XYZ2*)xyz2, this->xyz2DataNum);
        distGm.SetGlobalBuffer((__gm__ DTYPE_DIST*)dist, this->coreDataNum);
        indicesGm.SetGlobalBuffer((__gm__ DTYPE_INDICES*)indices, this->coreDataNum);

        // pipe.InitBuffer(inQueueXYZ1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_XYZ1));
        // pipe.InitBuffer(inQueueXYZ2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_XYZ2));
        // pipe.InitBuffer(outQueueDIST, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_DIST));
        // pipe.InitBuffer(outQueueINDICES, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_INDICES));

    }
    __aicore__ inline void Process() {
        float ux,uy,uz;
        float x,y,z;
        float best1, best2, best3;
        int besti1, besti2, besti3;
        for(uint32_t i=0; i<this->coreDataNum; i+=3)
        {
            ux = xyz1Gm.GetValue(i);
            uy = xyz1Gm.GetValue(i+1);
            uz = xyz1Gm.GetValue(i+2);
            for(j=0; j<this->xyz2DataNum; j+=3)
            {
                x = xyz2Gm.GetValue(j);
                y = xyz2Gm.GetValue(j+1);
                z = xyz2Gm.GetValue(j+2);
            }
        }
        // SetValue(count*M*N + i*N+j, xGm.GetValue(count*M*N + i*N+j));
    }
//     __aicore__ inline void Process() {

//         int32_t loopCount = this->tileNum;
//         this->processDataNum = this->tileDataNum;
//         for (int32_t i = 0; i < loopCount; i++) {
//             if (i == this->tileNum - 1) {
//               this->processDataNum = this->tailDataNum;
//             }
//             CopyIn(i);
//             Compute(i);
//             CopyOut(i);
//         }
//     }

// private:
//     __aicore__ inline void CopyIn(int32_t progress)
//     {
//         LocalTensor<DTYPE_XYZ1> xyz1Local = inQueueXYZ1.AllocTensor<DTYPE_XYZ1>();
//         LocalTensor<DTYPE_XYZ2> xyz2Local = inQueueXYZ2.AllocTensor<DTYPE_XYZ2>();

//         DataCopy(xyz1Local, xyz1Gm[progress * this->tileDataNum], this->processDataNum);
//         DataCopy(xyz2Local, xyz2Gm[progress * this->tileDataNum], this->processDataNum);

//         inQueueXYZ1.EnQue(xyz1Local);
//         inQueueXYZ2.EnQue(xyz2Local);
//     }
//     __aicore__ inline void Compute(int32_t progress)
//     {
//         LocalTensor<DTYPE_XYZ1> xyz1Local = inQueueXYZ1.DeQue<DTYPE_XYZ1>();
//         LocalTensor<DTYPE_XYZ2> xyz2Local = inQueueXYZ2.DeQue<DTYPE_XYZ2>();
//         LocalTensor<DTYPE_DIST> distLocal = outQueueDIST.AllocTensor<DTYPE_DIST>();
//         LocalTensor<DTYPE_INDICES> indicesLocal = outQueueINDICES.AllocTensor<DTYPE_INDICES>();

       
//         outQueueDIST.EnQue<DTYPE_DIST>(distLocal);
//         outQueueINDICES.EnQue<DTYPE_INDICES>(indicesLocal);
//         inQueueXYZ1.FreeTensor(xyz2Local);
//         inQueueXYZ2.FreeTensor(xyz2Local);
//     }
//     __aicore__ inline void CopyOut(int32_t progress)
//     {
//         LocalTensor<DTYPE_DIST> distLocal = outQueueDIST.DeQue<DTYPE_DIST>();
//         LocalTensor<DTYPE_INDICES> indicesLocal = outQueueINDICES.DeQue<DTYPE_INDICES>();
//         DataCopy(distGm[progress * this->tileDataNum], distLocal, this->processDataNum);
//         DataCopy(indicesGm[progress * this->tileDataNum], indicesLocal, this->processDataNum);
//         outQueueDIST.FreeTensor(distLocal);
//         outQueueINDICES.FreeTensor(indicesLocal);
//     }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmpBuffer, signbitBuffer;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueXYZ1, inQueueXYZ2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueDIST, outQueueINDICES;
    GlobalTensor<DTYPE_XYZ1> xyz1Gm;
    GlobalTensor<DTYPE_XYZ2> xyz2Gm;
    GlobalTensor<DTYPE_DIST> distGm;
    GlobalTensor<DTYPE_INDICES> indicesGm;
    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
    uint32_t xyz2DataNum;
    // LocalTensor<DTYPE_X> signbit;
};

extern "C" __global__ __aicore__ void three_nn(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelFastGeluGrad op;
    op.Init(xyz1, xyz2, dist, indices, 
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum, tiling_data.xyz2DataNum);  
    op.Process();
}
// extern "C" __global__ __aicore__ void three_nn(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices, GM_ADDR workspace, GM_ADDR tiling) {
//     GET_TILING_DATA(tiling_data, tiling);
//     // TODO: user kernel impl
// }